package com.tuniu.agents.agent;

import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.prompt.PromptManager;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.util.CollectionUtils;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * Implementation of Agent based on AgentOptions
 */
public class OptionsAgent extends Agent {

    /**
     * Agent configuration
     */
    protected AgentOptions agentOptions;
    /**
     * Request LLM Client
     */
    protected final ChatClient chatClient;
    /**
     * Chat memory
     */
    protected final ChatMemory chatMemory;
    /**
     * Prompt manager
     */
    protected final PromptManager promptManager;
    protected final ChatEnvironmentMemory chatEnvironmentMemory;
    protected final String agentId;

    public OptionsAgent(AgentManagerGroup agentManagerGroup, AgentOptions agentOptions) {
        this(agentManagerGroup.getModelBuilder(), agentManagerGroup.getChatMemory(), agentManagerGroup.getAdvisors(), agentManagerGroup.getPromptManager(), agentManagerGroup.getChatEnvironmentMemory(), agentOptions);
    }

    public OptionsAgent(ChatClient.Builder modelBuilder, ChatMemory chatMemory, List<Advisor> advisors, PromptManager promptManager, ChatEnvironmentMemory chatEnvironmentMemory, AgentOptions agentOptions) {
        this((agentOptions.getChatOptions() == null ? modelBuilder.clone() : modelBuilder.clone().defaultOptions(agentOptions.getChatOptions()))
                .defaultSystem(promptManager.getTemplateContent(agentOptions.getPromptTemplateId()))
                .defaultAdvisors(advisors)
                .defaultTools(CollectionUtils.isEmpty(agentOptions.getTools()) ? new String[0] : agentOptions.getTools().toArray(new String[0]))
                .build(), chatMemory, promptManager, chatEnvironmentMemory, agentOptions);
    }

    public OptionsAgent(ChatClient chatClient, ChatMemory chatMemory, PromptManager promptManager, ChatEnvironmentMemory chatEnvironmentMemory, AgentOptions agentOptions) {
        this.chatClient = chatClient;
        this.chatMemory = chatMemory;
        this.promptManager = promptManager;
        this.chatEnvironmentMemory = chatEnvironmentMemory;
        this.agentId = agentOptions.getAgentId();
        this.agentOptions = agentOptions;
    }

    public AgentOptions getAgentOptions() {
        return agentOptions;
    }

    @Override
    public String getAgentId() {
        return this.agentId;
    }

    @Override
    public Flux<ChatResponse> streamResponse(String conversationId, List<Message> messages, Map<String, Object> context) {
        return chatClientRequestSpec(conversationId, messages, context)
                .stream()
                .chatResponse().map(this::wrapper);
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String, Object> context) {
        return chatClientRequestSpec(conversationId, messages, context)
                .stream()
                .content();
    }

    @Override
    public ChatResponse callResponse(String conversationId, List<Message> messages, Map<String, Object> context) {
        return this.wrapper(chatClientRequestSpec(conversationId, messages, context)
                .call()
                .chatResponse());
    }

    @Override
    public String call(String conversationId, List<Message> messages, Map<String, Object> context) {
        return chatClientRequestSpec(conversationId, messages, context)
                .call()
                .content();
    }

    public ChatClient.ChatClientRequestSpec chatClientRequestSpec(String conversationId, List<Message> messages) {
        return chatClientRequestSpec(conversationId, messages, null);
    }

    @Override
    public ChatClient.ChatClientRequestSpec chatClientRequestSpec(String conversationId, List<Message> messages, Map<String,Object> context) {
        return chatClient.prompt()
                .system(s -> s.params(chatEnvironmentMemory.get(conversationId)))
                .messages(messages)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId).
                        param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, agentOptions.advisorParam.get(CHAT_MEMORY_RETRIEVE_SIZE_KEY)).
                        param(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY, this.agentId));
    }

}